package com.xwq.xwqaiagent.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;

import java.util.List;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

/**
 *ai恋爱大师向量存储与转换（云PgVectorStore）
 */
@Configuration
public class PgVectorVectorStoreConfig {

  @Resource
  private LoveAppDocumentLoader loveAppDocumentLoader;

  @Resource
  private MyTokenTextSplitter myTokenTextSplitter;

  @Resource
  private MyKeywordEnricher myKeywordEnricher;

  @Bean
  public VectorStore pgVectorVectorStore(
      @Qualifier("secondaryJdbcTemplate") JdbcTemplate pgJdbcTemplate,
      EmbeddingModel dashscopeEmbeddingModel) {

    PgVectorStore vectorStore = PgVectorStore.builder(pgJdbcTemplate, dashscopeEmbeddingModel)
            .dimensions(1536)
            .distanceType(COSINE_DISTANCE)
            .indexType(HNSW)
            .initializeSchema(true)
            .schemaName("public")
            .vectorTableName("vector_store")
            .maxDocumentBatchSize(10000)
            .build();
    // 加载文档
    List<Document> documents = loveAppDocumentLoader.loadMarkdowns();

//    //自主切分文档
//    List<Document> splitCustomized = myTokenTextSplitter.splitCustomized(documents);
//    //添加到向量数据库
//    vectorStore.add(splitCustomized);

    //调用基于ai的元信息增强器
    List<Document> enrichDocuments = myKeywordEnricher.enrichDocuments(documents);

    vectorStore.add(enrichDocuments);
   return vectorStore;
  }
}
